from emergency_plan.generate import generate_plan
from emergency_plan.query_rewrite import query_rewrite, query_rewrite_retrieval_only
from emergency_plan.recall import recall
from emergency_plan.rerank import rerank


class EmergencyPlanPipeline:
    def __init__(self, event_type, event_subtype, optimize_field=None, optimize_suggest=None, history=None, retrieval_only=False):
        self.event_type = event_type
        self.event_subtype = event_subtype
        self.optimize_field = optimize_field
        self.optimize_suggest = optimize_suggest
        self.history = history or []
        self.retrieval_only = retrieval_only  # 新增参数：是否使用纯检索模式

    async def process(self):
        # 构建查询
        query = f"事件类型：{self.event_type}，事件子类：{self.event_subtype}"
        
        print(f"🔍 EmergencyPlanPipeline开始处理")
        print(f"🔍 查询: {query}")
        print(f"🔍 优化字段: {self.optimize_field}")
        print(f"🔍 优化建议: {self.optimize_suggest}")
        print(f"🔍 历史对话数量: {len(self.history) if self.history else 0}")
        print(f"🔍 处理模式: {'纯检索模式' if self.retrieval_only else 'RAG模式'}")
        
        try:
            # 查询重写
            print(f"🔍 开始查询重写...")
            if self.retrieval_only:
                # 使用纯检索模式
                rewrite_result = await query_rewrite_retrieval_only(query)
            else:
                # 使用RAG模式（大模型）
                rewrite_result = await query_rewrite(query)
            print(f"✅ 查询重写完成: {rewrite_result}")
            
            # 召回相关文档
            print(f"🔍 开始文档召回...")
            recall_result = recall(rewrite_result)
            print(f"✅ 文档召回完成，召回数量: {len(recall_result) if recall_result else 0}")
            
            # 重排序
            print(f"🔍 开始文档重排序...")
            rerank_result = rerank(recall_result)
            print(f"✅ 文档重排序完成，重排序数量: {len(rerank_result) if rerank_result else 0}")
            
            if self.retrieval_only:
                # 纯检索模式：直接返回检索到的文档
                print(f"🔍 纯检索模式：直接返回检索到的相关文档")
                
                # 构建返回结果
                result = {
                    "status": "success",
                    "message": "检索完成，返回相关文档",
                    "mode": "retrieval_only",
                    "data": {
                        "event_type": self.event_type,
                        "event_subtype": self.event_subtype,
                        "optimize_field": self.optimize_field,
                        "optimize_suggest": self.optimize_suggest,
                        "related_documents": []
                    }
                }
                
                # 将检索到的文档转换为可序列化的格式
                for i, article in enumerate(rerank_result):
                    doc_info = {
                        "rank": i + 1,
                        "title": article.name,
                        "content": article.content,
                        "file_id": article.file_id,
                        "score": article.score,
                        "permission_type": article.permission_type,
                        "permission_set": article.permission_set
                    }
                    result["data"]["related_documents"].append(doc_info)
                
                print(f"✅ 应急预案检索完成，返回 {len(result['data']['related_documents'])} 个相关文档")
                return result
            else:
                # RAG模式：使用大模型生成应急预案
                print(f"🔍 RAG模式：开始生成应急预案...")
                result = await generate_plan(
                    rerank_result, 
                    self.event_type, 
                    self.event_subtype, 
                    self.optimize_field, 
                    self.optimize_suggest, 
                    self.history
                )
                print(f"✅ 应急预案生成完成")
                return result
            
        except Exception as e:
            print(f"❌ EmergencyPlanPipeline处理失败: {e}")
            import traceback
            traceback.print_exc()
            raise e 